from absl import app, flags
import datetime
import math
import typing

import pandas as pd
from pathlib import Path

from collections import defaultdict
from coin.base.format_util import (float_fmt, percentage_fmt)
from coin.exchange.kr_rest.product.product_impl import (
    generate_product_from_str2)
from coin.proto.coin_query_pb2 import AccountBalance
from coin.support.accounting.logic.account_proto_util import (
    convert_to_balance_dict,
    sum_balances)
from coin.support.pta.app.coinfof_pta_dumper import CoinFOFPtaDumper
from coin.tool.strat_monitor.util.monitor_util import load_pnl_adjust_config
from coin.util.protobuf.util import protos_to_df
from xunkemgmt_client.client.util.query_util import (
    query_accounts, query_exchange_rates, query_balance_snapshots)
from xunkemgmt_client.tool.slack_noti import send_to_slack
from xunkemgmt_client.support.accounting.logic.constants import (
    FIAT_LIST, STABLE_COIN_LIST)
from xunkemgmt_client.client.api_client import XunkemgmtClient
from xunkemgmt_client.base.datetime_util import to_timestamp_int
import coin2.service.account.account_info_service_pb2 as ais_pb2


BALANCE_BUSINESS_UNIT_LIST = ('Coin', 'CoinFOF', 'Cold', 'Day1mm', 'DigitalFund', 'Prex')


FLAGS = flags.FLAGS


def _format_df(df):
  for column in df:
    if column[1] == 'PCT':
      fmt_func = percentage_fmt
    else:
      fmt_func = float_fmt
    df[column] = df[column].apply(fmt_func)
  return df


def _format_balance(df):
  formatted_df = pd.DataFrame()
  for coin_type in df.index.unique(level='coin_type'):
    coin_df = df.loc[[coin_type]]
    total = coin_df.loc[:, pd.IndexSlice[:, ['USD', 'PCT']]].sum()
    coin_df = coin_df.append(pd.Series(name=(coin_type, 'Total')))
    coin_df.loc[(coin_type, 'Total')] = total
    formatted_df = pd.concat([formatted_df, coin_df])
  return _format_df(formatted_df)


def _format_summary(df):
  df = _format_df(df)
  formatted = pd.DataFrame()
  columns = list(df.columns.get_level_values(0))
  for column in columns:
    formatted[column] = df[column].apply(
        lambda x: '%s (%s)' % (x['USD'], x['PCT']), axis=1)
  return formatted


def _get_coinfof_adjusted_balance():
  pta_dumper = CoinFOFPtaDumper()
  trading_date = datetime.datetime.utcnow().date()
  td_str = trading_date.strftime('%Y-%m-%d')
  coinfof_pnl_summary = pta_dumper.get_pnl_report(trading_date, trading_date)
  total_reserve = pd.Series()
  accm_pnl_in_usd = 0
  for strat_name, summary in coinfof_pnl_summary.items():
    summary = summary[td_str]
    reserve = summary['reserve'].dropna()
    total_reserve = total_reserve.add(reserve, fill_value=0)
    diff_est = \
        (summary['balance_est'].fillna(0) - summary['reserve_est'].fillna(0)).sum()
    diff_est = diff_est if diff_est > 0 else 0
    accm_pnl_in_usd += diff_est
  adjusted_coinfof_balance = total_reserve.add(
      pd.Series([accm_pnl_in_usd], index=['USD']), fill_value=0)
  return adjusted_coinfof_balance


def _create_balance_dataframe(response_list):
  balance_list = [
      {} if response is None else
      {balance.currency: balance.total for balance in response.each_balance
          if not math.isclose(balance.total, 0, abs_tol=1e-5)}
      for response in response_list]
  return pd.DataFrame(balance_list)


def _create_all_balance_df(response_list: typing.List[AccountBalance],
                           usd_price_dict: typing.Dict[str, float]
) -> pd.DataFrame:
  if len(response_list) == 0: return
  df = _create_balance_dataframe(response_list)
  new_price_dict = {}
  for key in usd_price_dict:
    new_price_dict[key.split('-')[0]] = usd_price_dict[key]
  estimates_in_usd = []
  for _, row in df.iterrows():
    estimate = 0
    for col in df.columns:
      if col in new_price_dict and new_price_dict[col] and pd.notnull(row[col]):
        estimate += float(row[col] * new_price_dict[col])
    estimates_in_usd.append(estimate)
  balance_info_df = protos_to_df(response_list)
  balance_info_df = balance_info_df[['market_type', 'exchange', 'owner']]
  df_response = pd.concat([
      balance_info_df,
      pd.Series(estimates_in_usd, name='estimates_in_usd'),
      df,
  ], axis=1)
  df_response.loc['Total'] = df_response.sum(axis=0, numeric_only=True)
  df = df_response.sort_values([
    "market_type", "exchange", "owner"])
  return df


def _dump_account_balance_estimate(balance_list, output_path=None):
  Path(output_path).parent.mkdir(parents=True, exist_ok=True)
  balance_list.to_csv(output_path, encoding='utf-8')


def _load_balance_sheet_config(config_path, trading_date=None):
  if config_path is None:
    return pd.DataFrame()
  config_dict = load_pnl_adjust_config(config_path)
  config_list = [dict(**config, trading_date=date)
    for date, configs in config_dict.items() for config in configs]
  config_df = pd.DataFrame(config_list)
  config_df['trading_date'] = pd.to_datetime(config_df['trading_date'])
  if trading_date:
    config_df = config_df[config_df['trading_date'] <= pd.to_datetime(trading_date)]
  assert len(config_df) > 0
  latest_ids = config_df.groupby(['business_unit', 'symbol'])['trading_date'].idxmax()
  config_df = config_df.iloc[latest_ids]
  config_df = config_df.drop(columns=['trading_date'])
  return config_df


def _parse_options_product(options_df):
  tmp_df = options_df.copy()
  for idx in range(len(tmp_df)):
    product = generate_product_from_str2(
      'Options', 'Otc', 'v1', tmp_df.at[idx, 'symbol'])
    tmp_df.at[idx, 'strike_price'] = product.strike_price
    tmp_df.at[idx, 'symbol'] = product.base.currency
  return tmp_df


def _get_debt_subject(trading_date):
  debt_df = _load_balance_sheet_config(
    FLAGS.balance_debt_config, trading_date)
  debt_df.rename({'amount': 'debt'}, inplace=True, axis=1)
  debt_df['cur_price'] = query_exchange_rates(
    debt_df['symbol'], 'USD', trading_date=trading_date)
  debt_df['debt_usd'] = debt_df['debt'] * debt_df['cur_price']
  return debt_df


def _get_options_payoff_subject(trading_date):
  options_df = _load_balance_sheet_config(
    FLAGS.options_asset_config, trading_date)
  options_df = _parse_options_product(options_df)
  options_df.rename({'amount': 'option'}, inplace=True, axis=1)
  options_df['cur_price'] = query_exchange_rates(
    options_df['symbol'], 'USD', trading_date=trading_date)
  options_df['due'] = options_df.apply(
    lambda x: x['option'] * min(x['strike_price'], x['cur_price']), axis=1)
  options_df['options_payoff'] = options_df['option'] * \
    options_df['cur_price'] - options_df['due']
  return options_df


def _check_no_price(ticker_dict: typing.Dict[str, float],
                    exclude_config: typing.Optional[str]=None,
                    trading_date: typing.Optional[datetime.datetime]=None
) -> typing.Dict[str, float]:
  missing_ticker_list = [key for key in ticker_dict if ticker_dict[key] is None]
  if exclude_config:
    exclude_df = pd.read_csv(exclude_config)
    exclude_list = exclude_df['asset'].to_list()
    missing_ticker_list = [ticker for ticker in missing_ticker_list if ticker not in exclude_list]
  if missing_ticker_list and trading_date is None:
    new_missing_list = []
    kline_prices = query_exchange_rates(
      missing_ticker_list, 'USD',
      datetime.datetime.utcnow().date() - datetime.timedelta(days=1)
    )
    for i in range(len(missing_ticker_list)):
      if kline_prices[i] is not None:
        currency = missing_ticker_list[i]
        ticker_dict[currency] = kline_prices[i]
      else:
        new_missing_list.append(missing_ticker_list[i])
    missing_ticker_list = new_missing_list
  assert len(missing_ticker_list) == 0, f"no price for {','.join(missing_ticker_list)}"
  return ticker_dict


def create_balance_bu_summary(debt_df, options_df, balance_df):
  debt_df['debt'] = debt_df['debt_usd']
  debt_df = debt_df.groupby(['business_unit']).sum()
  options_df = options_df.groupby(['business_unit']).sum()
  debt_df.index = pd.MultiIndex.from_product([debt_df.index, ['USD']])
  options_df.index = pd.MultiIndex.from_product([options_df.index, ['USD']])
  balance_df['asset'] = balance_df.sum(axis=1)
  balance_df = pd.concat([balance_df, debt_df, options_df], axis=1).fillna(0)
  balance_df['equity'] = balance_df['asset'] - balance_df['debt'] + balance_df['options_payoff']
  balance_df = balance_df[['asset', 'options_payoff', 'debt', 'equity']]
  balance_df.columns = pd.MultiIndex.from_product([balance_df.columns, ['USD']])
  for column in list(balance_df.columns.get_level_values(0)):
    balance_df[column, 'PCT'] = balance_df[column, 'USD'] / balance_df[column, 'USD'].sum()
  balance_df = balance_df.groupby(level=0).sum()
  balance_df.loc['TOTAL'] = balance_df.sum()
  balance_df.index = balance_df.index.str.upper()
  return balance_df


def get_all_balance_df(account_balances: typing.List[AccountBalance],
                       trading_date: typing.Optional[datetime.date]
) -> pd.DataFrame:
  all_total_balance = convert_to_balance_dict(sum_balances(account_balances))
  coin_list = list(all_total_balance.keys())
  usd_price_list = query_exchange_rates(coin_list, 'USD', trading_date=trading_date)
  usd_price_dict = {'%s-USD' % coin: price for coin, price in
                    zip(coin_list, usd_price_list)}
  balance_df = _create_all_balance_df(account_balances, usd_price_dict)
  return balance_df


def get_balance_per_bu_df(account_balances: typing.List[AccountBalance]
) -> pd.DataFrame:
  account_list = query_accounts(as_proto=True)
  acct_key = lambda x: (x.market_type, x.exchange, x.owner)
  account_bu_dict = {acct_key(account): account.business_unit
    for account in account_list}
  balance_bu_dict = defaultdict(list)
  for acct_balance in account_balances:
    biz_unit = account_bu_dict[acct_key(acct_balance)].lower()
    balance_bu_dict[biz_unit].append(acct_balance)
  bu_total_balances = [convert_to_balance_dict(sum_balances(balance_bu))
    for balance_bu in balance_bu_dict.values()]
  balance_df = pd.DataFrame(bu_total_balances,
                            index=balance_bu_dict.keys())
  return balance_df


def query_balance(business_units: typing.List[str],
                  trading_date: typing.Optional[datetime.date]
) -> typing.List[AccountBalance]:
  balance_list = []
  if trading_date is None:
    data = query_balance_snapshots(business_units=business_units,
                                   as_proto=True)
  else:
    with XunkemgmtClient() as client:
      data = client.query_balance_interval_histories(
        ais_pb2.QueryBalanceIntervalHistoriesRequestProto(
          start_timestamp=to_timestamp_int(trading_date),
          end_timestamp=to_timestamp_int(
          trading_date + datetime.timedelta(1)),
          business_units=business_units)
      ).histories
  for acct_balance in data:
    if acct_balance:
      balance_info = acct_balance.balance_info
      balance_list.append(AccountBalance(
        market_type=balance_info.market_type,
        exchange=balance_info.exchange,
        owner=balance_info.owner,
        each_balance=acct_balance.each_balance
      ))
  return balance_list


def check_balance_report(balance_acct_df,
                         slack_receiver,
                         cross_check_file_path,
                         diff_threshold=1e-3):
  total_estimate_usd = balance_acct_df.loc['Total', 'estimates_in_usd']
  msg_output = ''
  prev_balance = pd.read_csv(cross_check_file_path)
  prev_total_estimate_usd = prev_balance.loc[prev_balance.index[-1], 'estimates_in_usd']
  if abs(total_estimate_usd - prev_total_estimate_usd) / max(
      total_estimate_usd, prev_total_estimate_usd) > diff_threshold:
    diff_detail = balance_acct_df.loc['Total', :] - prev_balance.iloc[-1, :]
    diff_detail = diff_detail[abs(diff_detail /
      balance_acct_df.loc['Total', :]) > diff_threshold]
    diff_detail = diff_detail.T.dropna()
    diff_detail = diff_detail.to_frame(name='RAW')
    diff_detail = diff_detail.drop(index='estimates_in_usd')
    msg_output = '```' + 'UTC 0 am balance report invalid:'
    msg_output += f'\nTotal diff estimate in usd: {total_estimate_usd - prev_total_estimate_usd}'
    msg_output += '\nBalance diff details: \n' + diff_detail.to_string(float_format=float_fmt) + '```'
  else:
    msg_output = '```UTC 0 am balance report confirmed```'
  send_to_slack(msg_output, slack_receiver, 'msg')


def get_balance_summary() -> typing.Dict:
  business_units = [bu.lower() for bu in BALANCE_BUSINESS_UNIT_LIST]
  trading_date = datetime.datetime.strptime(FLAGS.trading_date, '%Y%m%d').date() \
                 if FLAGS.trading_date is not None else None
  account_balances = query_balance(business_units, trading_date)
  all_balance_df = get_all_balance_df(account_balances, trading_date)
  balance_per_bu_df = get_balance_per_bu_df(account_balances)

  coinfof_adj_balance_df = pd.DataFrame(
      _get_coinfof_adjusted_balance(), columns=['coinfof_adj']).T
  debt_df = _get_debt_subject(trading_date)
  options_df = _get_options_payoff_subject(trading_date)
  currencies = list(set(list(balance_per_bu_df.columns) +
    list(debt_df['symbol']) + list(options_df['symbol'])))
  prices = query_exchange_rates(currencies, 'USD', trading_date=trading_date)
  if FLAGS.check_no_price:
    prices_dict = _check_no_price(dict(zip(currencies, prices)),
      FLAGS.exclude_estimate_config, trading_date=trading_date)
    currencies, prices = prices_dict.keys(), prices_dict.values()
  price_df = pd.DataFrame(
      [{currency: price for currency, price in zip(currencies, prices)}],
      index=['price'])
  debt_ccy_df = debt_df.groupby('symbol').sum()
  options_ccy_df = options_df.groupby('symbol').sum()

  balance_per_bu_df = pd.concat(
      [price_df, balance_per_bu_df, coinfof_adj_balance_df], sort=True).T
  asset, asset_adj = pd.Series(), pd.Series()
  for bu in business_units:
    asset = asset.add(balance_per_bu_df[bu], fill_value=0)
    if bu == 'coinfof':
      bu = 'coinfof_adj'
    asset_adj = asset_adj.add(balance_per_bu_df[bu], fill_value=0)
  balance_per_bu_df['asset'] = asset
  balance_per_bu_df['asset_adj'] = asset_adj
  balance_per_bu_df.columns = pd.MultiIndex.from_product(
      [balance_per_bu_df.columns, ['RAW']])
  balance_per_bu_df.index.name = 'currency'
  columns = list(balance_per_bu_df.columns.get_level_values(0))
  for column in columns:
    balance_per_bu_df[column, 'USD'] = balance_per_bu_df[column, 'RAW'] * \
                                        balance_per_bu_df['price', 'RAW']
  balance_per_bu_df['options_payoff', 'USD'] = options_ccy_df['options_payoff']
  balance_per_bu_df['debt', 'USD'] = debt_ccy_df['debt_usd']
  balance_per_bu_df['equity', 'USD'] = balance_per_bu_df['asset', 'USD'].subtract(
    balance_per_bu_df['debt', 'USD'], fill_value=0).add(
      balance_per_bu_df['options_payoff', 'USD'], fill_value=0)
  balance_per_bu_df['equity_adj', 'USD'] = \
    balance_per_bu_df['asset_adj', 'USD'].subtract(
    balance_per_bu_df['debt', 'USD'], fill_value=0).add(
    balance_per_bu_df['options_payoff', 'USD'], fill_value=0)
  columns += ['options_payoff', 'debt', 'equity', 'equity_adj']
  for column in columns:
    balance_per_bu_df[column, 'PCT'] = \
      balance_per_bu_df[column, 'USD'] / balance_per_bu_df[column, 'USD'].sum()
  balance_bu_df = balance_per_bu_df.loc[:, pd.IndexSlice[business_units, 'USD']].T
  summary_bu_df = create_balance_bu_summary(debt_df, options_df, balance_bu_df)

  balance_per_bu_df['coin_type'] = [
      'FIAT' if currency in FIAT_LIST + STABLE_COIN_LIST else
      'COIN' for currency in balance_per_bu_df.index]
  balance_per_bu_df = balance_per_bu_df.reset_index().set_index(
      ['coin_type', 'currency'])
  estimate_df = balance_per_bu_df[
    ['asset', 'asset_adj', 'options_payoff', 'debt', 'equity', 'equity_adj']].fillna(0)
  estimate_df = estimate_df.loc[:, pd.IndexSlice[:, ['USD', 'PCT']]]
  summary_df = pd.concat([
      estimate_df.loc[('COIN', 'BTC')].to_frame(name='BTC').T,
      estimate_df.groupby(level=0).sum()])
  summary_df.loc['COIN'] = summary_df.loc['COIN'] - summary_df.loc['BTC']
  summary_df.rename(index={'COIN': 'ALT'}, inplace=True)
  summary_df.loc['TOTAL'] = summary_df.sum()

  summaries_dict = {
    'debt_df': debt_df,
    'options_df': options_df,
    'balance_per_bu_df': balance_per_bu_df,
    'balance_per_acct_df': all_balance_df,
    'balance_sheet_per_coin_df': summary_df,
    'balance_sheet_per_bu_df': summary_bu_df,
  }
  return summaries_dict


def print_balance_summary(summaries_dict, slack_receiver, cross_check_file_path=None):
  print("\n\n\n")
  print("Sing", datetime.datetime.now())
  print("UTC", datetime.datetime.utcnow())

  balance_per_bu_df = _format_balance(summaries_dict['balance_per_bu_df'])
  balance_sheet_per_coin_df = _format_summary(summaries_dict['balance_sheet_per_coin_df'])
  balance_sheet_per_bu_df = _format_summary(summaries_dict['balance_sheet_per_bu_df'])

  print("\n==================================\n")
  print("Debt detail:\n")
  print(summaries_dict['debt_df'].to_string(float_format=float_fmt))
  print("\nOptions detail:\n")
  print(summaries_dict['options_df'].to_string(float_format=float_fmt))
  print("\nTotal:\n")
  print(balance_per_bu_df.to_string())
  print("\nSummary:\n")
  print(balance_sheet_per_coin_df.to_string())
  print(balance_sheet_per_bu_df.to_string())

  balance_per_acct_df = summaries_dict['balance_per_acct_df']
  dt = datetime.datetime.utcnow().strftime("%Y%m%d%H")
  if FLAGS.trading_date:
    csv_file_path = f'/tmp/{dt}_check_balance_report.csv'
  else:
    csv_file_path = f'/tmp/{dt}_balance_report.csv'
  _dump_account_balance_estimate(balance_per_acct_df, csv_file_path)

  if slack_receiver is not None:
    send_to_slack(balance_per_bu_df['asset'].to_string(), slack_receiver, 'file')
    for df in (balance_sheet_per_coin_df, balance_sheet_per_bu_df):
      msg = '```' + df.to_string() + '```'
      send_to_slack(msg, slack_receiver, 'msg')
    with open(csv_file_path) as f:
      send_to_slack(f.read(), slack_receiver, 'file')
    if cross_check_file_path:
      check_balance_report(balance_per_acct_df, slack_receiver, cross_check_file_path)


def main(_):
  pd.set_option('display.width', 1000)
  pd.set_option('display.max_rows', None)
  pd.set_option('display.max_columns', None)
  summaries_dict = get_balance_summary()
  if FLAGS.print_ratio:
    print_balance_summary(summaries_dict,
                          FLAGS.slack_receiver,
                          cross_check_file_path=FLAGS.cross_check_file_path)


if __name__ == "__main__":
  flags.DEFINE_string('account_group', None, '')
  flags.DEFINE_string('trading_date', None, '%Y%m%d')
  flags.DEFINE_boolean('print_ratio', False, '')
  flags.DEFINE_string('slack_receiver', None, 'slack_receiver')
  flags.DEFINE_string('cross_check_file_path', None, 'cross_check_file_path')
  flags.DEFINE_boolean('check_no_price',
                       False,
                       'raise exception when price is NA.')
  flags.DEFINE_string('exclude_estimate_config',
                      "../../coin_deploy/support_accounting/balance_config/estimate_exclude_list.csv",
                      'path of price exclude config.')
  flags.DEFINE_string(
      'balance_debt_config',
      "../../coin_deploy/support_accounting/balance_config/balance_debt_config.json", '')
  flags.DEFINE_string(
      'options_asset_config',
      "../../coin_deploy/support_accounting/balance_config/options_asset_config.json", '')
  app.run(main)
